
import os
import sys

sys.path.append(os.path.split(sys.path[0])[0])
import shutil
from time import time
import SimpleITK as sitk
import scipy.ndimage as ndimage
from multiprocessing.dummy import Pool
import math
import numpy as np
'''
这里把遵义乳腺癌的MRI图像的前96层，分为3个32层
'''

data_root = r"E:\AllData\zunYiBreast"  # 数据的根路径，可以设置为“nii_path”的上一级目录

volume_path = r"E:\AllData\zunYiBreast\nii\volume"
seg_path = r"E:\AllData\zunYiBreast\nii\segmentation"

pretreatment_volume_path = os.path.join(data_root, 'cutCenter256-3/volume')
pretreatment_seg_path = os.path.join(data_root, 'cutCenter256-3/segmentation')

def cutCenter(volume_array):
    target_size = 256
    shape_x = 448
    shape_y = 448
    x_begin = shape_x // 2 - target_size // 2
    x_end = shape_x // 2 + target_size // 2
    y_begin = shape_y // 2 - target_size // 2
    y_end = shape_y // 2 + target_size // 2

    volume_array = volume_array[:, x_begin:x_end, y_begin:y_end]
    # print("one successful!",file)
    return volume_array

def cutTo32(file:str,processfile_path,save_path):
    sample_duration=32
    sample_size=256
    # 将CT和标签入读内存
    volume = sitk.ReadImage(os.path.join(processfile_path, file), sitk.sitkInt16)
    volume_array = sitk.GetArrayFromImage(volume)
    Spacing=volume.GetSpacing()
    Direction=volume.GetDirection()
    Origin=volume.GetOrigin()

    volume_array1=cutCenter(volume_array[:32])
    volume_array2=cutCenter(volume_array[32:64])
    volume_array3 = cutCenter(volume_array[64:96])

    volume1 = sitk.GetImageFromArray(volume_array1)
    volume2 = sitk.GetImageFromArray(volume_array2)
    volume3 = sitk.GetImageFromArray(volume_array3)

    volume1.SetDirection(Direction)
    volume1.SetOrigin(Origin)
    volume1.SetSpacing(Spacing)
    volume2.SetDirection(Direction)
    volume2.SetOrigin(Origin)
    volume2.SetSpacing(Spacing)
    volume3.SetDirection(Direction)
    volume3.SetOrigin(Origin)
    volume3.SetSpacing(Spacing)

    print(volume_array1.shape,volume_array2.shape,volume_array3.shape)
    print( os.path.join(save_path, file.split('.')[0]+'(1-32)'+'.nii.gz'))


    sitk.WriteImage(volume1, os.path.join(save_path, file.split('.')[0]+'(1-32)'+'.nii.gz'))
    sitk.WriteImage(volume2, os.path.join(save_path, file.split('.')[0] + '(32-64)' + '.nii.gz'))
    sitk.WriteImage(volume3, os.path.join(save_path, file.split('.')[0] + '(64-96)' + '.nii.gz'))
    # print("one successful!",file)



if __name__ == '__main__':

    # if os.path.exists(para.pretreatment_path):
    #     shutil.rmtree(para.pretreatment_path)
    if not os.path.exists(pretreatment_volume_path):
        os.makedirs(pretreatment_volume_path)
    if not os.path.exists(pretreatment_seg_path):
        os.makedirs(pretreatment_seg_path)
    start = time()
    file_list=os.listdir(volume_path)
    print("处理MRI")
    for file in file_list:
        cutTo32(file,processfile_path = volume_path,save_path = pretreatment_volume_path)

    file_list = os.listdir(seg_path)
    print("处理标注")
    for file in file_list:
        cutTo32(file, processfile_path = seg_path, save_path = pretreatment_seg_path)



